#!/usr/bin/env python

#$Id: flatten.py 86 2009-05-27 20:54:29Z dkmagee $
# -------------------------------------------------------------
#
__version__      = '$Revision: 86 $ '[11:-3]
__version_date__ = '$Date: 2009-05-27 20:54:29 +0000 (Wed, 27 May 2009) $ '[7:-3]
__author__       = 'R. Bouwens, <bouwens@ucolick.org>, D. Magee, <magee@ucolick.org>'

import numpy as N
import ndimage as nd
import pyfits,os,imagestats,glob,sys
import shutil

def measurenoise(flatdata,dq,staticmask,clip):
    ndq = (dq == 0) * (staticmask == 1) * (flatdata < clip) * (flatdata > -clip*1.3)
    hclip = clip / 2
    flatdata_med = nd.median_filter(flatdata.astype(N.float32),size=3)
    ndq *= (flatdata_med < hclip)
    dqd = (ndq == 1)
    good = flatdata[dqd]
    stats = imagestats.ImageStats(good.astype(N.float32), fields='stddev')
    return stats.stddev

def measurenoiseFN(FN,staticmask):
    data = pyfits.getdata(FN)
    dq = pyfits.getdata(FN,ext='DQ')
    staticmask = pyfits.getdata(staticmask)    
    noise = measurenoise(data,dq,staticmask,0.03)
    return noise

def findbkg(data,flatdata,dq,staticmask,clip,size,goodfact=0.25):
    ndq = (dq == 0) * (staticmask == 1) * (flatdata < clip) * (flatdata > -clip*1.3)
    hclip = clip / 2
    flatdata_med = nd.median_filter(flatdata.astype(N.float32),size=3)
    ndq *= (flatdata_med < hclip)

    (xdim,ydim) = data.shape
    ngridx = xdim / size
    ngridy = ydim / size

    Median = N.zeros((ngridx,ngridy), dtype=N.float64)
    NContr = N.zeros((ngridx,ngridy), dtype=N.int32)
    for i in range(ngridx):
        for j in range(ngridy):
            xl = i*size
            xh = N.clip(xl+size,0,xdim)
            yl = j*size
            yh = N.clip(yl+size,0,ydim)
            datasl = data[xl:xh,yl:yh]
            dqsl = ndq[xl:xh,yl:yh]
            good = (dqsl == 1)
            curdata = datasl[good]
            NContr[i,j] = len(curdata)
            if NContr[i,j] > 2:
                stats = imagestats.ImageStats(curdata.astype(N.float32), fields='median',binwidth=0.00005)
                Median[i,j] = stats.median
            else:
                Median[i,j] = 0
    MedianM = N.zeros((ngridx+1,ngridy+1), dtype=N.float64)
    for i in range(ngridx):
        for j in range(ngridy):
            notdone = 1
            sizem = 1
            while notdone:
                xl = N.clip(i-sizem,0,ngridx)
                xh = N.clip(i+sizem+1,0,ngridx)
                yl = N.clip(j-sizem,0,ngridy)
                yh = N.clip(j+sizem+1,0,ngridy)  
                datasl = Median[xl:xh,yl:yh]
                dqsl = NContr[xl:xh,yl:yh]
                good = (dqsl > size*size*goodfact)
                curdata = datasl[good]
                if len(curdata) > 1:
                    stats = imagestats.ImageStats(curdata.astype(N.float32), fields='median',binwidth=0.00005)
                    if not N.isnan(stats.median):
                        notdone = 0
                if notdone:
                    sizem += 1

            MedianM[i+1,j+1] = stats.median
            if i == 0:
                MedianM[0,j+1] = stats.median
            if j == 0:
                MedianM[i+1,0] = stats.median
                if i == 0:
                    MedianM[0,0] = stats.median

    newmedian = nd.affine_transform(MedianM.astype(N.float32),[[1./size,0],[0,1./size]],output_shape=(256,256),offset=(0.5,0.5))
    newmedian *= staticmask
    return newmedian

def findbkgquad(data,flatdata,dq,staticmask,clip,size,goodfact=0.5):
  bkg = N.zeros((256,256), dtype=N.float32)
  for m in range(2):
    for n in range(2):
      xl = m*128
      xh = m*128+128
      yl = n*128
      yh = n*128+128
      cbkg = findbkg(data[yl:yh,xl:xh],data[yl:yh,xl:xh],dq[yl:yh,xl:xh],staticmask[yl:yh,xl:xh],0.01,size=9)
      bkg[yl:yh,xl:xh] = cbkg
  return bkg

def derivestep(data,flatdata,dq,staticmask,clip):
    ndq = (dq == 0) * (staticmask == 1) * (flatdata < clip) * (flatdata > -clip*1.3)
    hclip = clip / 2
    flatdata_med = nd.median_filter(flatdata.astype(N.float32),size=3)
    ndq *= (flatdata_med < hclip)
    ndq1 = ndq[15:127,1:127].ravel()
    ndq2 = ndq[15:127,129:256].ravel()
    ndq3 = ndq[129:241,1:127].ravel()
    ndq4 = ndq[129:241,129:241].ravel()
    perc1 = len(N.where(ndq1 == 1)[0])/float(len(ndq1))
    perc2 = len(N.where(ndq2 == 1)[0])/float(len(ndq2))
    perc3 = len(N.where(ndq3 == 1)[0])/float(len(ndq3))
    perc4 = len(N.where(ndq4 == 1)[0])/float(len(ndq4))
    L = [(perc1,1),(perc2,2),(perc3,3),(perc4,4)]
    L.sort()
    (perc,i) = L[0]
    if perc < 0.45:
        L = []
        for j in range(44):
	  if j < 24:
	    if i % 2 == 1:
              x1 = 1 + 5 * j
              x2 = 12 + 5 * j            
	    else:
              x1 = 129 + 5 * j
              x2 = 140 + 5 * j
            med1 = data[112:127,x1:x2]
            med2 = data[129:142,x1:x2]
            dqsl1 = ndq[112:127,x1:x2]
            dqsl2 = ndq[129:142,x1:x2]
            good1 = (dqsl1 == 1)
            good2 = (dqsl2 == 1)
            curdata1 = med1[good1]
            curdata2 = med2[good2]
	    if (N.sum(good1.ravel().astype(N.int32)) > 100) and (N.sum(good2.ravel().astype(N.int32)) > 100):
              stats1 = imagestats.ImageStats(curdata1.astype(N.float32), fields='median',binwidth=0.00005)
              stats2 = imagestats.ImageStats(curdata2.astype(N.float32), fields='median',binwidth=0.00005)
 	      if i < 3:
	        diff = stats2.median - stats1.median
	      else:
	        diff = stats1.median - stats2.median
	      print 'x1,x2 = ',x1,x2,' diff = ',diff
              L.append(diff)
	  else:
	    if i < 3:
              y1 = 15 + 5 * (j-24)
              y2 = 26 + 5 * (j-24)
	    else:	
              y1 = 129 + 5 * (j-24)
              y2 = 140 + 5 * (j-24)         
            med1 = data[y1:y2,112:127]
            med2 = data[y1:y2,129:142]
            dqsl1 = ndq[y1:y2,112:127]
            dqsl2 = ndq[y1:y2,129:142]
            good1 = (dqsl1 == 1)
            good2 = (dqsl2 == 1)
            curdata1 = med1[good1]
            curdata2 = med2[good2]
	    if (N.sum(good1.ravel().astype(N.int32)) > 100) and (N.sum(good2.ravel().astype(N.int32)) > 100):
              stats1 = imagestats.ImageStats(curdata1.astype(N.float32), fields='median',binwidth=0.00005)
              stats2 = imagestats.ImageStats(curdata2.astype(N.float32), fields='median',binwidth=0.00005)
              if i % 2 == 1:
  	        diff = stats2.median - stats1.median
	      else:
  	        diff = stats1.median - stats2.median
	      print 'y1,y2 = ',y1,y2,' diff = ',diff
              L.append(diff)
	if len(L) > 0:
          med = N.median(L)
          if not N.isnan(med):
            if i == 1:
              data[0:128,0:128] += med
            if i == 2:
              data[0:128,128:256] += med
            if i == 3:
              data[128:256,0:128] += med
            if i == 4:
              data[128:256,128:256] += med


def flatten(inname, outname, staticmask, brightobj):
    newf = pyfits.open(inname)
    data = newf['SCI'].data
    dq = pyfits.getdata(inname,ext='DQ')
    staticmask = pyfits.getdata(staticmask)
    cmedian = findbkg(data,data,dq,staticmask,0.08,size=11)
    flatdata = data - cmedian
    noise = measurenoise(flatdata,dq,staticmask,0.03)
    print 'RMS noise = ',noise
    if brightobj:
      derivestep(data,flatdata,dq,staticmask,noise*2)
    cmedian = findbkg(data,flatdata,dq,staticmask,noise*2,size=11)
    data -= cmedian
    newf.writeto(outname,clobber=True)
    newf.close()

